# -*- coding: utf-8 -*-
'''
Created on 17.10.2019

@author: yu03
'''
import glob
import re
import datetime
import cv2
import matplotlib.pyplot as plt
from scipy import signal
from FFT_Interpolation import FFT_interpolation_2, line_cal, line_cal_fix
import numpy as np
import os
from scipy.optimize import curve_fit
from mpl_toolkits.mplot3d import Axes3D
from Video_Unified import folder_path, video_names, hor_index, ver_index, hor_lines, ver_lines, line_num
import sys

# calculation_mode = 'diff' ### DC removed
calculation_mode = 'DC' ### DC Not removed


def fit_func(x, a, b, c):
    return a*(x-b)**2 + c

''' 
    General Parameters
'''
fs_cam = 50
Lamda = 633e-9
pix_size = 5.3e-6
V_x, V_y, V_z = 0, 0, 0


# for video in video_name:
#     i += 1
#     file_name = re.findall(r"Compare\\(.+).avi", video)[0]
#     np_result_name = folder_path + '\\' + file_name[0] + '_lines' + '.npy'
#     print(video)
#     print(i, file_name)

#     print(txt_name)


for video in video_names:
    ''' 
        File name define
    '''
    file_name = re.findall(r"6__40um_30s_openmode\\(.+).avi", video)[0]
    if calculation_mode == 'DC':
        np_result_name = folder_path + '\\results\\' + file_name + '_lines' + '.npy'
    elif calculation_mode == 'diff':
        np_result_name = folder_path + '\\results_diff\\' + file_name + '_lines_diff' + '.npy'
    else:
        sys.exit('FFT Mode Error:\n calculation_mode = %s'%calculation_mode)
    print('File:', file_name, '; ', datetime.datetime.now().time())
#     print(np_result_name)
    img_set = []
    
    ''' 
        Reading from Video file
    '''
    cap = cv2.VideoCapture(video)
    frame_total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
#     print(frame_total, "Frames")
#     print('ok')
    for k in range(frame_total):
        ret, frame = cap.read() 
        img = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        img_set.append(img)
#         cv2.imshow('frame',img)
#         if cv2.waitKey(1) & 0xFF == ord('q'):
#             break
#     cap.release()
#     cv2.destroyAllWindows()

    
    ''' 
        Finding Phase Spectrum Component
    '''
    hor_m_k_num_set = []
    ver_m_k_num_set = []
    for img in img_set[500:2000][::5]:
        hor_center = np.array(img[hor_index])
        ver_center = img[:, ver_index]
        
        if calculation_mode == 'diff':
            hor_center = np.diff(hor_center.astype('int'))
            hor_center = np.concatenate(([0], hor_center))
            ver_center = np.diff(ver_center.astype('int'))
            ver_center = np.concatenate(([0], ver_center))
        elif calculation_mode == 'DC':
            pass
        else:
            sys.exit('FFT Mode Error:\n calculation_mode = %s'%calculation_mode)
        
        hor_f_fit, hor_phase_estim, hor_m_k_num = line_cal(hor_center)
        ver_f_fit, ver_phase_estim, ver_m_k_num = line_cal(ver_center)
        
        hor_m_k_num_set.append(hor_m_k_num)
        ver_m_k_num_set.append(ver_m_k_num) 
    m_k_num_hor = int(np.average(hor_m_k_num_set))
    m_k_num_ver = int(np.average(ver_m_k_num_set))
    print(frame_total, "Frames;  ", 'Chosen Components:', m_k_num_hor, m_k_num_ver, '  ', datetime.datetime.now().time() )
    
    ''' 
        Processing
    '''    
    hor_f_set, ver_f_set = [], []
    hor_phi_set, ver_phi_set = [], []
    frame_num = 0
    for img in img_set:
    # for img in img_set[0:2]:
        frame_num += 1
        for i in range(line_num):
            hor_center = np.array(img[hor_lines[i]])
            ver_center = img[:, ver_lines[i]]
            
            if calculation_mode == 'diff':
                hor_center = np.diff(hor_center.astype('int'))
                hor_center = np.concatenate(([0], hor_center))
                ver_center = np.diff(ver_center.astype('int'))
                ver_center = np.concatenate(([0], ver_center))
            elif calculation_mode == 'DC':
                pass
            else:
                sys.exit('FFT Mode Error:\n calculation_mode = %s'%calculation_mode)
            
            hor_f_fit, hor_phase_estim, hor_m_k_num = line_cal_fix(hor_center, m_k_num_hor, m_k_num_ver)
            ver_f_fit, ver_phase_estim, ver_m_k_num = line_cal_fix(ver_center, m_k_num_hor, m_k_num_ver)
            
            hor_f_set.append(hor_f_fit)
            hor_phi_set.append(hor_phase_estim)
            ver_f_set.append(ver_f_fit)
            ver_phi_set.append(ver_phase_estim)
        if frame_num%10 == 0:
            print('Frame: ', frame_num, '   ',datetime.datetime.now().time())   
    hor_f_set = np.array(hor_f_set)
    ver_f_set = np.array(ver_f_set)
    hor_f_set = hor_f_set.reshape(line_num, frame_num, order='F')
    ver_f_set = ver_f_set.reshape(line_num, frame_num, order='F')
 
    hor_phi_set = np.array(hor_phi_set)
    ver_phi_set = np.array(ver_phi_set)
    hor_phi_set = hor_phi_set.reshape(line_num, frame_num, order='F')
    ver_phi_set = ver_phi_set.reshape(line_num, frame_num, order='F')
 
    hor_angle_set = (V_x-Lamda*hor_f_set/2)*1e6 ### urad
    ver_angle_set = (V_x-Lamda*ver_f_set/2)*1e6 ### urad
    hor_length_set = np.unwrap(hor_phi_set)/4/np.pi*Lamda*1e9 ### nm
    ver_length_set = np.unwrap(ver_phi_set)/4/np.pi*Lamda*1e9 ### nm
      
    print('line, frame:',np.shape(hor_angle_set))
    print('')
    f = open(np_result_name,'ab')
    np.save(f, hor_angle_set, allow_pickle=True)
    np.save(f, ver_angle_set, allow_pickle=True)
    np.save(f, hor_length_set, allow_pickle=True)
    np.save(f, ver_length_set, allow_pickle=True)
    
     

''' 
    Plotting
'''   
y = [np.arange(frame_num)]*line_num
y = np.array(y)
  
fig = plt.figure('hor_tilt')
x = [hor_lines]*frame_num
x = np.array(x).T
ax = fig.gca(projection='3d')
plt.gca().patch.set_facecolor('white')
for i in range(line_num):
    ax.plot(x[i,:],y[i,:],hor_angle_set[i,:])
   
fig = plt.figure('ver_tilt')
x = [ver_lines]*frame_num
x = np.array(x).T
ax = fig.gca(projection='3d')
plt.gca().patch.set_facecolor('white')
for i in range(line_num):
    ax.plot(x[i,:],y[i,:],ver_angle_set[i,:])
     
     
fig = plt.figure('hor_length')
x = [hor_lines]*frame_num
x = np.array(x).T
ax = fig.gca(projection='3d')
plt.gca().patch.set_facecolor('white')
for i in range(line_num):
    ax.plot(x[i,:],y[i,:],hor_length_set[i,:])
     
fig = plt.figure('ver_length')
x = [hor_lines]*frame_num
x = np.array(x).T
ax = fig.gca(projection='3d')
plt.gca().patch.set_facecolor('white')
for i in range(line_num):
    ax.plot(x[i,:],y[i,:],ver_length_set[i,:])
 
plt.show()